﻿using System;
using System.Collections.Generic;
using Algoritmia.DataStructures.Maps;
using Algoritmia.Utils;

namespace Algoritmia.Problems.Compression
{
    public class HuffmanCodeBuilder
    {
        private readonly Func<IMap<string, string>> _createMap;

        public HuffmanCodeBuilder(Func<IMap<string,string>> createMap = null)
        {
            _createMap = createMap ?? (() => AlgoritmiaFactory.Maps<string, string>.Default);
        }

        public IEnumerable<KeyValuePair<string,string>> BuildCode(HuffmanTree tree)
        {
            if ( tree == null)
            {
                throw new ArgumentNullException("tree");
            }

            var codes = _createMap();

            Action<HuffmanTree,string> traverse = null;
            traverse = (t, s) =>
                           {
                               if (t.Left == null && t.Right == null)
                               {
                                   codes[t.Symbol] = s;
                                   return;
                               }
                                traverse((HuffmanTree)t.Left, s + "0");
                                traverse((HuffmanTree)t.Right, s + "1");
                           };
            traverse(tree, "");

            return codes;
        }
    }
}
